{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ae8e2ae0-da25-44e7-ae82-024173150a15",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests \n",
    "\n",
    "docs_url = 'https://github.com/alexeygrigorev/llm-rag-workshop/raw/main/notebooks/documents.json'\n",
    "docs_response = requests.get(docs_url)\n",
    "documents_raw = docs_response.json()\n",
    "\n",
    "documents = []\n",
    "\n",
    "for course in documents_raw:\n",
    "    course_name = course['course']\n",
    "\n",
    "    for doc in course['documents']:\n",
    "        doc['course'] = course_name\n",
    "        documents.append(doc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "57de60e5-b96c-499c-a7cf-0f30fc33b324",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'text': \"Yes, even if you don't register, you're still eligible to submit the homeworks.\\nBe aware, however, that there will be deadlines for turning in the final projects. So don't leave everything for the last minute.\",\n",
       " 'section': 'General course-related questions',\n",
       " 'question': 'Course - Can I still join the course after the start date?',\n",
       " 'course': 'data-engineering-zoomcamp'}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "documents[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ce7d0d18-5c07-4010-9f90-bbd021f110c8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<minsearch.minsearch.Index at 0x7c43acd9fc20>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import minsearch\n",
    "\n",
    "index = minsearch.Index(\n",
    "    text_fields=[\"question\", \"text\", \"section\"],\n",
    "    keyword_fields=[\"course\"]\n",
    ")\n",
    "\n",
    "index.fit(documents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "aa755a08-b98d-4e92-8994-04e6108499d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "\n",
    "client = OpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b21237c3-80e9-429c-a089-d45428087046",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search(query):\n",
    "    boost = {'question': 3.0, 'section': 0.5}\n",
    "\n",
    "    results = index.search(\n",
    "        query=query,\n",
    "        filter_dict={'course': 'data-engineering-zoomcamp'},\n",
    "        boost_dict=boost,\n",
    "        num_results=5\n",
    "    )\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8cc5784e-6515-42e5-be62-8fb915df1088",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_prompt(query, search_results):\n",
    "    prompt_template = \"\"\"\n",
    "You're a course teaching assistant. Answer the QUESTION based on the CONTEXT from the FAQ database.\n",
    "Use only the facts from the CONTEXT when answering the QUESTION.\n",
    "\n",
    "QUESTION: {question}\n",
    "\n",
    "CONTEXT: \n",
    "{context}\n",
    "\"\"\".strip()\n",
    "\n",
    "    context = \"\"\n",
    "    \n",
    "    for doc in search_results:\n",
    "        context = context + f\"section: {doc['section']}\\nquestion: {doc['question']}\\nanswer: {doc['text']}\\n\\n\"\n",
    "    \n",
    "    prompt = prompt_template.format(question=query, context=context).strip()\n",
    "    return prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "97d35dec-c25f-472d-b961-20d5c30902ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "def llm(prompt):\n",
    "    response = client.chat.completions.create(\n",
    "        model='gpt-4o',\n",
    "        messages=[{\"role\": \"user\", \"content\": prompt}]\n",
    "    )\n",
    "    \n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "8602f40b-ad3b-49c9-b3cc-051a79c888bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rag(query):\n",
    "    search_results = search(query)\n",
    "    prompt = build_prompt(query, search_results)\n",
    "    answer = llm(prompt)\n",
    "    return answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "5fd4497b-c5d5-4258-b950-6b35d1af4ec5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'To run Kafka with Java, you need to execute the following command in your project directory:\\n\\n```sh\\njava -cp build/libs/<jar_name>-1.0-SNAPSHOT.jar:out src/main/java/org/example/JsonProducer.java\\n```\\n\\nIf you\\'re using Python and encounter the \"Module \\'kafka\\' not found\" error when trying to run `producer.py`, you should create a virtual environment, activate it, and install the required packages as per `requirements.txt`:\\n\\n1. Create a virtual environment (run only once):\\n   ```sh\\n   python -m venv env\\n   ```\\n\\n2. Activate the virtual environment:\\n   - On MacOS/Linux:\\n     ```sh\\n     source env/bin/activate\\n     ```\\n   - On Windows:\\n     ```sh\\n     env\\\\Scripts\\\\activate\\n     ```\\n\\n3. Install the necessary packages:\\n   ```sh\\n   pip install -r ../requirements.txt\\n   ```\\n\\nEnsure your Docker images are running if required. To deactivate the virtual environment when done, use:\\n```sh\\ndeactivate\\n```'"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rag('how do I run kafka?')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "385b012f-4905-422d-8d7c-3d542dfe5a7c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Yes, you can still enroll in the course even if it has already started. You are also eligible to submit the homework assignments. However, make sure to pay attention to the deadlines for turning in the final projects to avoid leaving everything until the last minute.'"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rag('the course has already started, can I still enroll?')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
